import os
import torch
import numpy as np
from torch.utils.data.dataset import Dataset
import cv2 as cv
from utils import letter_box


class MyDataset(Dataset):
    def __init__(self):
        super().__init__()
        img_prepath = 'E:/dataset/flower'
        classes = os.listdir(img_prepath)
        self.num_classes = len(classes)
        self.img_prepath = []
        self.img_num = np.zeros(self.num_classes).astype('int32')
        for i in range(self.num_classes):
            self.img_prepath.append(img_prepath + '/' + classes[i])
            img_names = os.listdir(self.img_prepath[i])
            self.img_num[i] = len(img_names)
        self.img_total = int(sum(self.img_num))

    def __len__(self):
        return self.img_total

    def item_remap(self, item):
        for i in range(self.num_classes):
            if item < self.img_num[i]:
                return i, item
            item -= self.img_num[i]

    def __getitem__(self, item):
        label, cnt = self.item_remap(item)
        img_prepath = self.img_prepath[label]
        img_path = img_prepath + '/' + os.listdir(img_prepath)[cnt]
        image = cv.imread(img_path)
        image = letter_box(image, 224)
        image = image[:, :, ::-1].transpose(2, 0, 1)/256
        return torch.tensor(image, dtype=torch.float32), torch.tensor(label, dtype=torch.int64)
